package com.feidee.fd.sml.algorithm.component.feature

import com.feidee.fdspark.transformer.{ColSelector, MetaStorage, ModelType}
import org.apache.spark.ml.feature.Imputer
import org.apache.spark.ml.{Pipeline, PipelineModel, PipelineStage}
import org.apache.spark.sql.DataFrame

/**
  * @Author: dongguosheng
  * @Date: 2019/3/18 18:52
  * @Review songhaicheng
  * @Email: guosheng_dong@sui.com
  */
case class ImputeEncoderParam(
                               override val input_pt: String,
                               override val output_pt: String,
                               override val hive_table: String,
                               override val flow_time: String,
                               override val inputCol: String,
                               override val outputCol: String,
                               override val preserveCols: String,
                               override val modelPath: String,
                               // 缺失值填充策略，支持 [mean, median]，默认 mean
                               strategy: String,
                               // 缺失值的定义，默认 Double.NaN
                               missingValue: Double
                             ) extends FeatureParam {

  def this() = this(null, null, null, null, "input", "features", null, null, "mean", Double.NaN)

  override def verify(): Unit = {
    super.verify()
    val strategies = Array("mean", "median")
    require(strategies.contains(strategy.toLowerCase), s"param strategy only accepts [${strategies.mkString(", ")}]," +
      s" but has $strategy")
  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("strategy" -> strategy)
    map += ("missingValue" -> missingValue)
    map
  }
}


class ImputeEncoder extends AbstractFeatureEncoder[ImputeEncoderParam] {
  override def setUp(param: ImputeEncoderParam, data: DataFrame): Array[PipelineStage] = {
    val imputer = new Imputer()
      .setInputCols(param.inputCol.split(","))
      .setOutputCols(param.outputCol.split(","))
      .setMissingValue(param.missingValue)
      .setStrategy(param.strategy)
    Array(imputer)
  }

  override def train(param: ImputeEncoderParam, data: DataFrame): PipelineModel = {
    // 把 preserveCols 取出来
    val preserveCols = if (tool.isNotNull(param.preserveCols)) {
      param.preserveCols.split(",")
    } else {
      // 不填 preserveCols 时，赋上一个空数组，不然上面会返回一个长度为 1 内容空字符串的数组！
      new Array[String](0)
    }

    // 结果列
    val cols = preserveCols ++ param.outputCol.split(",")

    // 存储训练参数元信息
    val meta = new MetaStorage()
      .setModelType(ModelType.Feature)
      .setParameters(param.toMap)
      .setFields(data.schema.fieldNames)

    // 最终结果只保留结果列
    val colSelector = new ColSelector()
      .setCols(cols)

    // 执行流程：存储元信息 》 特征组件处理过程 》 保留结果列
    val pipeline = new Pipeline().setStages(Array(meta) ++ setUp(param, data) ++ Array(colSelector))
    val model = pipeline.fit(data)

    model
  }
}


object ImputeEncoder {

  def apply(paramStr: String): Unit = {
    new ImputeEncoder()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    ImputeEncoder(args(0))
  }
}